import numpy as np
import librosa
import os
from pandas import DataFrame
import pandas as pd
import matplotlib.pyplot as plt


data_directory = "F:/heart_data/2022_challenge_new/the-circor-digiscope-phonocardiogram-dataset-1.0.3/training_data"
tsv_path = os.path.join(data_directory,'50690_TV.tsv')
wav_path = os.path.join(data_directory,'50690_TV.wav')
# print(pd.read_csv(tsv_path, delimiter='\t'))
def tsv_load(tsvname):
    """读取tsv文件内容,不需要close函数"""
    with open(tsvname, "r") as f:
        txt_data = f.read()
    head = ["start", "end", "period"]
    data = txt_data.split("\n")[:-1]
    # 遍历每一行
    for l in data:
        sgmt = l.split("\t")
        if sgmt[2] != "0":
            head = np.vstack([head, sgmt])
    return head[1:]

def get_beat_time(data_tsv):
    beat_time = []
    S1_time = []
    S2_time = []
    beat_start = 0
    for i in range(len(data_tsv)):
        if data_tsv[i][2] == '1':
            S1_start = float(data_tsv[i][0])
            S1_end = float(data_tsv[i][1])
            beat_start = float(data_tsv[i][0])
        if data_tsv[i][2] == '4' and beat_start != 0:
            beat_end = float(data_tsv[i][1])
            S2_start = float(data_tsv[i][0])
            S2_end = float(data_tsv[i][1])
            beat_time.append([beat_start, beat_end])
            S1_time.append([S1_start, S1_end])
            S2_time.append([S2_start, S2_end])
    return beat_time,S1_time,S2_time

#
# data_tsv = tsv_load(tsv_path)
# beat_time, S1_time ,S2_time= get_beat_time(data_tsv)
# recording, fs = librosa.load(wav_path, sr=4000)  # 分割（3s不重叠）
# cut = list()
# for num in range(len(beat_time)):  # 将每个片段写入对应的听诊区文件夹
#     start = int(beat_time[num][0]*fs)
#     end = int(beat_time[num][1]*fs)
#     S1_start = int(S1_time[num][0]*fs)
#     S1_end = int(S1_time[num][1]*fs)
#     S1_recording = recording[S1_start:S1_end]
#     S2_start = int(S2_time[num][0] * fs)
#     S2_end = int(S2_time[num][1] * fs)
#     S2_recording = recording[S2_start:S2_end]
#     max_S=max(S1_recording.max(),S2_recording.max())
#     small = recording[start:end]
#     small_pre = small / max_S
#     cut.append(small)



# #统计每个PCG记录的心拍最大时长，平均时长
# beats_len=['max','mean']
# for f in sorted(os.listdir(data_directory)):
#     root, extension = os.path.splitext(f)
#     beats_num=0
#     start_time='nan'
#     current_beat_len=[]
#     if extension == '.tsv':
#         path_tsv = os.path.join(data_directory,f)
#         data_tsv= tsv_load(path_tsv)
#         for i in range(len(data_tsv)):
#             if data_tsv[i][2] == '1':
#                 start_time=data_tsv[i][0]
#             if data_tsv[i][2] == '4' and start_time!='nan':
#                 end_time = data_tsv[i][1]
#                 beat_num = float(end_time)-float(start_time)
#                 current_beat_len.append(beat_num)
#                 beats_num += 1
#         max_len=np.max(current_beat_len)
#         mean_len=np.mean(current_beat_len)
#         a = [max_len,mean_len]
#         beats_len = np.vstack([beats_len,a])
#         beats = beats_len[1:]





